package com.rankeiot.core.anno.aspect;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.google.common.util.concurrent.RateLimiter;
import com.rankeiot.core.anno.RateLimit;
import jakarta.servlet.http.HttpServletRequest;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;


import java.util.concurrent.TimeUnit;

@Aspect
public class RateLimitAspect extends AbstractExpressionAspect {
    private Cache<String, RateLimiter> limiters = Caffeine.newBuilder().expireAfterAccess(10, TimeUnit.MINUTES).build();//new ConcurrentHashMap<>();

    @Before("@annotation(limit)")
    public void before(JoinPoint point, RateLimit limit) throws Throwable {
        StandardEvaluationContext ctx = new StandardEvaluationContext();
        ctx.registerFunction("getIp",RateLimitAspect.class.getDeclaredMethod("getIp"));
        String v = executeExpression(point, ctx, limit.value());
        RateLimiter limiter = limiters.get(v,k->RateLimiter.create(limit.permitsPerSecond()));
        if (!limiter.tryAcquire()) {
            throwBusinessException(point,limit.message());
        }
    }
    public static String getIp(){
        String ip = getIpAddr(getRequest());
        return ip;
    }
    private static HttpServletRequest getRequest()
    {
        return getRequestAttributes().getRequest();
    }
    private static ServletRequestAttributes getRequestAttributes()
    {
        RequestAttributes attributes = RequestContextHolder.getRequestAttributes();

        return (ServletRequestAttributes) attributes;
    }
    private static String getIpAddr(HttpServletRequest request)
    {
        if (request == null)
        {
            return "unknown";
        }
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip))
        {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip))
        {
            ip = request.getHeader("X-Forwarded-For");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip))
        {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip))
        {
            ip = request.getHeader("X-Real-IP");
        }

        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip))
        {
            ip = request.getRemoteAddr();
        }

        return "0:0:0:0:0:0:0:1".equals(ip) ? "127.0.0.1" : getMultistageReverseProxyIp(ip);
    }
    private static String getMultistageReverseProxyIp(String ip)
    {
        // 多级反向代理检测
        if (ip != null && ip.indexOf(",") > 0)
        {
            final String[] ips = ip.trim().split(",");
            for (String subIp : ips)
            {
                if (false == isUnknown(subIp))
                {
                    ip = subIp;
                    break;
                }
            }
        }
        return ip;
    }
    private static boolean isUnknown(String checkString)
    {
        return !StringUtils.hasText(checkString) || "unknown".equalsIgnoreCase(checkString);
    }
}
